use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::io;
use tokio::net::{TcpListener, TcpStream};
use tracing::info;
use tracing::log::warn;
use tracing_core::LevelFilter;
use tracing_subscriber::fmt::format::FmtSpan;
use tracing_subscriber::fmt::Layer;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

#[derive(Serialize, Deserialize, Debug)]
struct Config {
    listener_addr: String,
    upstream_addr: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    //proxy client traffic to upstream

    let console = Layer::new()
        .with_ansi(true)
        .with_line_number(true)
        .with_target(true)
        .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
        .with_thread_ids(true) // include the thread ID of the current thread
        .with_thread_names(true) // include the name of the current thread
        .json();

    let config = resolve_config();
    let config = Arc::new(config);

    tracing_subscriber::registry()
        .with(LevelFilter::INFO)
        .with(console)
        .init();

    info!("UpStream:{}", config.upstream_addr);
    info!("Listener:{}", config.listener_addr);
    let listener = TcpListener::bind(&config.listener_addr).await?;
    loop {
        let (client, addr) = listener.accept().await?;
        let cloned_config = Arc::clone(&config);
        info!("Incoming connection from {addr}");
        tokio::spawn(async move {
            let upstream = TcpStream::connect(&cloned_config.upstream_addr).await?;
            //proxy
            proxy(client, upstream).await?;
            Ok::<(), anyhow::Error>(())
        });
    }
    #[allow(unreachable_code)]
    Ok(())
}

async fn proxy(mut client: TcpStream, mut upstream: TcpStream) -> Result<()> {
    let (mut client_read, mut client_write) = client.split();
    let (mut upstream_read, mut upstream_write) = upstream.split();
    let client_to_upstream = io::copy(&mut client_read, &mut upstream_write);
    let upstream_to_client = io::copy(&mut upstream_read, &mut client_write);
    match tokio::try_join!(client_to_upstream, upstream_to_client) {
        Ok((n, m)) => {
            info!("proxied {n} bytes from client to upstream, {m} bytes from upstream to client")
        }
        Err(e) => warn!("error proxying: {e:?}"),
    }
    Ok(())
}

fn resolve_config() -> Config {
    Config {
        listener_addr: "0.0.0.0:8081".to_string(),
        upstream_addr: "0.0.0.0:8080".to_string(),
    }
}
